import numpy as np
import matplotlib.pyplot as plt
from scipy.integrate import quad


# 定义高斯窗口函数 g(t - τ, θ)
def gaussian_window(t, tau, theta):
    return (1 / (theta * np.sqrt(2 * np.pi))) * np.exp(-((t - tau) ** 2) / (2 * theta ** 2))

# 定义短时傅里叶变换 S(t, f)
def stft(t, f, x_t, theta):
    # 被积函数
    def integrand(tau):
        g_value = gaussian_window(t, tau, theta)
        return g_value * np.exp(-1j * 2 * np.pi * f * tau)

    # 执行从 -∞ 到 ∞ 的积分
    integral_value, _ = quad(integrand, -np.inf, np.inf)
    return x_t * integral_value  # 注意这里的 x(t) 保持不变


# 示例：定义输入信号 x(t)，例如简单的正弦波信号
def input_signal(t):
    return np.sin(2 * np.pi * 3 * t) # 1Hz 的正弦波信号


# 设置时间和频率的范围
time_range = np.linspace(0, 2, 100)  # 时间从0到2秒
frequency_range = np.linspace(1, 5, 50)  # 频率从1到10Hz
theta = 1 / frequency_range[0]  # 设置 θ 为第一个频率的倒数

# 创建一个空矩阵来存储时频变换结果
stft_matrix = np.zeros((len(frequency_range), len(time_range)), dtype=np.complex_)

# 计算每个时间和频率下的 S(t, f)
for i, t in enumerate(time_range):
    x_t = input_signal(t)  # 计算 x(t)
    for j, f in enumerate(frequency_range):
        theta = 1 / f  # 根据频率计算 θ
        stft_matrix[j, i] = stft(t, f, x_t, theta)

# 可视化：绘制时频图
plt.figure(figsize=(10, 6))
plt.pcolormesh(time_range, frequency_range, np.abs(stft_matrix), shading='auto', cmap='inferno')
plt.colorbar(label='Magnitude')
plt.xlabel('Time (s)')
plt.ylabel('Frequency (Hz)')
plt.title('Time-Frequency Representation (STFT with Gaussian Window)')
plt.show()